from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name="addxp",    # 注册到python的库名;经验证，要跟CUDAExtension里那个名字一样，不然会出错，CUDAExtension的名字是生成.so的名字
    include_dirs=["./"],
    ext_modules=[
        CUDAExtension(
            "addxp",
            ["pybind.cpp", "add.cu"],
        )
    ],
    cmdclass={
        "build_ext": BuildExtension
    }
)